package com.changdu.seatunnel.admin.service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.ListConsumerGroupOffsetsResult;
import org.apache.kafka.clients.admin.ListOffsetsResult;
import org.apache.kafka.clients.admin.NewPartitions;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.admin.OffsetSpec;
import org.apache.kafka.clients.admin.TopicDescription;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;

import com.changdu.seatunnel.admin.config.common.BusinessException;
import com.changdu.seatunnel.admin.config.enums.SyncErrorStatus;
import com.changdu.seatunnel.admin.pojo.TableRouter;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.google.common.collect.Sets;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class KafkaCreater {

	public static final String KAFKA_SERVERS = "bootstrap.servers";
	public static final String KAFKA_PROTOCOL = "security.protocol";
	public static final String KAFKA_SASL_MECHANISM = "sasl.mechanism";
	public static final String KAFKA_SASL_JAAS = "sasl.jaas.config";
	public static final String KAFKA_TOPICS = "topics";
	
	/**
	 * 创建 topic 及 partition
	 * @param kafkaConfig
	 */
	@SuppressWarnings("unchecked")
	public void createTopicIfNotExists(JsonNode nodeConf) {
		Properties kafkaConfig = buildKafkaProperties(nodeConf);
		
		List<TableRouter> tableRouters = (List<TableRouter>) kafkaConfig.get(KafkaCreater.KAFKA_TOPICS);
		if (tableRouters.isEmpty()) return;
		
		Map<String, Integer> topics = new HashMap<>();	
		tableRouters.stream().collect(Collectors.groupingBy(TableRouter::getTopic)).forEach((topic, routers) -> {
			int partition = routers.stream().mapToInt(TableRouter::getPartition).max().orElse(2);
			topics.put(topic, partition);
		});
		
		try (AdminClient adminClient = AdminClient.create(kafkaConfig)) {
			// exists topics
			Set<String> existsTopics = adminClient.listTopics().names().get().stream().filter(t -> topics.containsKey(t)).collect(Collectors.toSet());
			Map<String, TopicDescription> topicsDesc = adminClient.describeTopics(existsTopics).all().get();	
			Map<String, Integer> existsTopicInfos = topicsDesc.values().stream()
					.collect(Collectors.toMap(key -> key.name(), value -> value.partitions().stream().mapToInt(t -> t.partition()).max().orElse(2)));
			
			// 如果存在 topic, 则创建指定分区
			for (Entry<String, Integer> existsTopicInfo : existsTopicInfos.entrySet()) {
				String topicName = existsTopicInfo.getKey();
				Integer currentPartitionNum = existsTopicInfo.getValue();
				if (topics.get(topicName) > currentPartitionNum) {
					adminClient.createPartitions(Map.of(topicName, NewPartitions.increaseTo(topics.get(topicName) + 1))).all().get();
				}
			}
			
			int replicas = topicsDesc.values().stream().mapToInt(t -> t.partitions().stream().mapToInt(p -> p.replicas().size()).min().orElse(1)).min().orElse(1);
			// 如果不存在topic，则创建topic
			List<NewTopic> noExistTopics = Sets.newHashSet(topics.keySet()).stream()
					.filter(t -> !existsTopics.contains(t))
					.map(t -> new NewTopic(t, topics.get(t) + 1, (short) replicas))
					.collect(Collectors.toList());
			if (!noExistTopics.isEmpty()) {
				adminClient.createTopics(noExistTopics).all().get();
			}
		} catch (Exception e) {
			throw new BusinessException(SyncErrorStatus.KAFKA_REQUEST_EXCEPTION.getCode(), e.getMessage());
		}
	}

	/**
	 * 获取 kafka 消费 偏移量
	 * @param source
	 * @return
	 * @throws ExecutionException 
	 * @throws InterruptedException 
	 */
	public Long getLagFromTopic(JsonNode source) {
		Properties kafkaConfig = buildKafkaProperties(source);
		String consumerGroupId = source.get("consumer.group").asText("");
		
		try (AdminClient adminClient = AdminClient.create(kafkaConfig)) {
			ListConsumerGroupOffsetsResult offsetsResult = adminClient.listConsumerGroupOffsets(consumerGroupId);
			Map<TopicPartition, OffsetAndMetadata> consumerOffsets = offsetsResult.partitionsToOffsetAndMetadata().get();
			if (consumerOffsets.isEmpty()) return 0L;
			
			Map<TopicPartition, OffsetSpec> partitionOffsets = consumerOffsets.keySet().stream().collect(Collectors.toMap(tp -> tp, tp -> OffsetSpec.latest()));
			Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo> topicPartitionListOffsetsResultInfoMap = adminClient.listOffsets(partitionOffsets).all().get();
			
			
			return partitionOffsets.keySet().stream().mapToLong(partition -> {
				OffsetAndMetadata consumerOffset = consumerOffsets.get(partition);
				ListOffsetsResult.ListOffsetsResultInfo logEndOffsetInfo = topicPartitionListOffsetsResultInfoMap.get(partition);
				// 计算 Lag = log-end-offset - consumerOffset
				return logEndOffsetInfo.offset() - consumerOffset.offset();
			}).sum();
		} catch (Exception e) {
			log.error("get lag from topic error: {}", e.getMessage(), e);
		}
		return 0L;
	}
	
	/**
	 * 从配置模板中构建 kakfa 参数
	 * @param config
	 * @return
	 */
	protected Properties buildKafkaProperties(JsonNode config) {
		Properties properties = new Properties();
		JsonNode kafka = config.get("kafka.config");
		
		properties.put(KafkaCreater.KAFKA_SERVERS, config.get(KafkaCreater.KAFKA_SERVERS).asText());
		if (kafka.has(KafkaCreater.KAFKA_PROTOCOL)) {
			properties.put(KafkaCreater.KAFKA_PROTOCOL, kafka.get(KafkaCreater.KAFKA_PROTOCOL).asText());
		}
		if (kafka.has(KafkaCreater.KAFKA_SASL_MECHANISM)) {
			properties.put(KafkaCreater.KAFKA_SASL_MECHANISM, kafka.get(KafkaCreater.KAFKA_SASL_MECHANISM).asText());
		}
		if (kafka.has(KafkaCreater.KAFKA_SASL_JAAS)) {
			properties.put(KafkaCreater.KAFKA_SASL_JAAS, kafka.get(KafkaCreater.KAFKA_SASL_JAAS).asText());
		}
		List<TableRouter> tableRouters = new ArrayList<>();
		if (config.has("reroute")) {
			ArrayNode reroutes = (ArrayNode) config.get("reroute");
			for (JsonNode route : reroutes) {
				String pattern = route.get("pattern").asText("");
				String topic = route.get("topic").asText("");
				int partition = route.get("partition").asInt(0);
				tableRouters.add(new TableRouter(pattern, topic, partition));
			}
		}
		
		if (config.has("topic")) {
			String topic = config.get("topic").asText();
			Integer partition = config.has("partition") ? config.get("partition").asInt() : null;
			tableRouters.add(new TableRouter(null, topic, partition));
		}
		
		properties.put(KafkaCreater.KAFKA_TOPICS, tableRouters);
		return properties;
	}

}
